#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Author  : lei.wang
from app.util.model_provider.enums import ModelTypeEnum
from app.util.model_provider.huoshan.llm import VolcenginesChatCompletion
from app.util.model_provider.tongyi.llm import TongYiChatCompletion
from app import config
from app.core.pkg.log import logger


def get_model_provider(model_name: ModelTypeEnum):

    model_providers = {
        ModelTypeEnum.tongyi: TongYiChatCompletion(),
        ModelTypeEnum.huoshan: VolcenginesChatCompletion()
    }
    return model_providers.get(model_name, None)


if config.get('llmModel'):
    mapping = {
        'tongyi': ModelTypeEnum.tongyi,
        'huoshan': ModelTypeEnum.huoshan,
    }
    llm_obj = get_model_provider(mapping[config.get('llmModel').get('enable')])
    logger.success(f"加载模型生效, [{config.get('llmModel').get('enable')}]")
else:
    logger.warning('没有配置对应的大模型信息, 默认生效[tongyi]')
    llm_obj = get_model_provider(ModelTypeEnum.tongyi)
